package com.peng.leetcode.dynamic;

import java.util.HashMap;
import java.util.Map;

/**
 * Rob
 * <p>
 * 337. 打家劫舍 III
 *
 * @author: lupeng6
 * @create: 2021/2/18 17:11
 */
public class Rob2 {

    public static void main(String[] args) {
        TreeNode node1 = new TreeNode(3);
        TreeNode node2 = new TreeNode(4);
        TreeNode node3 = new TreeNode(5);
        TreeNode node4 = new TreeNode(1);
        TreeNode node5 = new TreeNode(3);
        TreeNode node6 = new TreeNode(1);


        node2.left = node4;
        node2.right = node5;
        node1.left = node2;

        node3.right = node6;

        node1.right = node3;

        System.out.println(new Rob2().rob(node1));
    }

    Map<TreeNode, Integer> select = new HashMap<>();
    Map<TreeNode, Integer> unSelect = new HashMap<>();


    public int rob(TreeNode root) {
        if (root == null) {
            return 0;
        }
        dfs(root);
        return Math.max(select.get(root), unSelect.get(root));
    }

    private void dfs(TreeNode node) {
        if (node == null) {
            return;
        }
        dfs(node.left);
        dfs(node.right);

        // 选中当前节点
        select.put(node, node.val + unSelect.getOrDefault(node.left, 0) + unSelect.getOrDefault(node.right, 0));

        // 未选中当前节点
        // 左节点 选中和不选中最大 + 有节点选中和不选中最大
        int leftMax = Math.max(select.getOrDefault(node.left, 0), unSelect.getOrDefault(node.left, 0));
        int rightMax = Math.max(select.getOrDefault(node.right, 0), unSelect.getOrDefault(node.right, 0));
        unSelect.put(node, leftMax + rightMax);
    }


    static class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;

        TreeNode(int x) {
            val = x;
        }

        @Override
        public String toString() {
            return "TreeNode{" +
                    "val=" + val +
                    '}';
        }
    }

}
